#!/usr/bin/env python
# -*- coding: utf-8 -*-
"""
An attempt at some simple, self-contained pygame-based examples.
Example 01
In short:
One static body: a big polygon to represent the ground
One dynamic body: a rotated big polygon
And some drawing code to get you going.
kne
"""
import random

import cv2
import pygame
# Box2D.b2 maps Box2D.b2Vec2 to vec2 (and so on)
from Box2D import b2FixtureDef, b2PolygonShape
from Box2D.b2 import world, dynamicBody
import numpy as np

import utils as utils
from Agent import Agent

# --- constants ---
# Box2D deals with meters, but we want to display pixels,
# so define a conversion factor:
TARGET_FPS = 60
SCREEN_WIDTH = 256
SCREEN_HEIGHT = 256
screen = pygame.display.set_mode((SCREEN_WIDTH, SCREEN_HEIGHT), 0, 32)
pygame.display.set_caption('Learning')
clock = pygame.time.Clock()
world = world(gravity=(0, 0), doSleep=True)
colors = {
    dynamicBody: (255, 0, 0, 255),
}
attachment = world.CreateDynamicBody(
    position=(5, 3),
    fixtures=b2FixtureDef(shape=b2PolygonShape(box=(0.5, 2)), density=2.0),
)
platform = world.CreateDynamicBody(
    position=(15, 5),
    fixtures=b2FixtureDef(shape=b2PolygonShape(box=(4, 0.5)), density=2, friction=0.6)
)
world.CreatePrismaticJoint(
    bodyA=attachment,
    bodyB=platform,
    anchor=(0, 5),
    axis=(1, 0),
    maxMotorForce=1000,
    enableMotor=True,
    lowerTranslation=-10,
    upperTranslation=10,
    enableLimit=True
)
bodiesInTheWorld = {
    'attachment': attachment,
    'platform': platform
}

worldContext = {
    'PPM': 10.0,  # pixels per meter
    'TARGET_FPS': TARGET_FPS,
    'TIME_STEP': 1.0 / TARGET_FPS,
    'SCREEN_WIDTH': 256,
    'SCREEN_HEIGHT': 256,
    'world': world,
    'screen': screen,
    'clock': clock,
    'colors': colors,
    'bodiesInTheWorld': bodiesInTheWorld
}
running = True
agent = Agent((worldContext['SCREEN_HEIGHT'], worldContext['SCREEN_WIDTH']))
utils.drawTheWorld(worldContext)
targetBboxTopLeft = np.array([160, 40])
while running:
    # agent starts to work
    # observedImage = cv2.imread('outputImages/temp.png')
    # agent.createCurrentScene(observedImage)
    # # after observation, check whether the target has reached
    # agent.matchObjectBetweenScene()
    # isDone = agent.checkTargetFininsed(1, targetBboxTopLeft)
    # if isDone:
    #     print('target finished')
    #     break
    # # agent.printScene(agent.currentScene)
    # agent.learnCauseEffect()
    # act = agent.getAction(1, targetBboxTopLeft)
    # print('performed action is:', act)
    utils.changeTheWorldByAction(act, worldContext)
    utils.drawTheWorld(worldContext)
pygame.quit()
print('Done!')
